'''
Created on Mar 10, 2013

@author: arenduchintala
'''
import math
import sys, os, re, locale
from pprint import pprint


unigramCounts = {}
bigramCounts = {}
trigramCounts = {}
emissionCounts = {}
seenObservations = {}
emissions = {}
trigramTransitions = {}
bigramTransitions = {}
unigramTransitions = {}

pi = {}
# s2i = {'I-GENE':2, 'O':3, '*':1, 'STOP':0}
s2i = {}
i2s = {}

def populateStateMap(_state):
    if (not s2i.has_key(_state)):
       s2i[_state] = _state  # len(s2i)
    return s2i[_state]
                
def MLEstimates(countFilePath):
    populateStateMap('*')
    populateStateMap('STOP')
    countStar = 0
    countStarStar = 0

    lines = open(countFilePath, 'r').readlines()
    for line in lines:
        line = line.rstrip()
        # #print line
        items = line.strip().split(" ")
        if (items[1].strip() == 'WORDTAG'):
            # 8 WORDTAG O localizes
            emissionCounts[(populateStateMap(items[2]) , items[3])] = int(items[0])
        
        elif (items[1].strip() == '1-GRAM'):
            # 749 1-GRAM I-GENE            
            unigramCounts[populateStateMap(items[2])] = int(items[0])

        elif (items[1].strip() == '2-GRAM'):
            # 749 2-GRAM * I-GENE
            bigramCounts[(populateStateMap(items[2]) , populateStateMap(items[3]))] = int(items[0])
            if populateStateMap(items[2]) == s2i['*'] :
                countStar = countStar + int(items[0])
                
        elif (items[1].strip() == '3-GRAM'):
            # 1813 3-GRAM I-GENE O STOP
            trigramCounts[(populateStateMap(items[2]) , populateStateMap(items[3]), populateStateMap(items[4]))] = int(items[0])
            if populateStateMap(items[2]) == s2i['*'] and populateStateMap(items[3]) == s2i['*']:
                countStarStar = countStarStar + int(items[0])

    # estimate emission probabilities
    unigramCounts[s2i['*']] = countStar;
    bigramCounts[(s2i['*'], s2i['*'])] = countStarStar;
    for k, v in emissionCounts.items():
        state = k[0]
        obs = k[1]
        key = (obs , state)
        seenObservations[obs] = 1
        value = math.log(v / float(unigramCounts[state]))
        # emissions[key] = value
        emissions[(obs, state)] = value
    
    # estimate transition probabilities
    sumUnigrams = sum(unigramCounts.values())
    for k, v in unigramCounts.items():
        # k = D 
        # count of D over all unigrams
        value = math.log(v / float(sumUnigrams))
        # probability_key = D/all unigrams
        probability_key = (k)
        unigramTransitions[probability_key] = value
        
    for k, v in bigramCounts.items():
        # k = (D,N)
        # count of D,N over count of D
        value = math.log(v / float(unigramCounts[k[0]]))
        probability_key = (k[1] , k[0])  # N/D
        bigramTransitions[probability_key] = value
        
    for k, v in trigramCounts.items():
        # k = (D,N,V) 
        state_old = (k[0] , k[1])
        # count of D,N,V over count of D,N
        value = math.log(v / float(bigramCounts[state_old]))
        probability_key = (k[2] , k[0] , k[1])  # V/D,N
        trigramTransitions[probability_key] = value
    

def getWordPossibleTags(word):    
    states = []
    for state in unigramCounts.keys():
        if emissions.has_key((word , state)):
            states.append(state)
    return states

def replaceRareWord(token):
    try:
        locale.atof(token)
        return '_RARE_NUMERIC_'
    except:
        if  len(re.findall(r'[A-Z]', token)) == len(token):
            return '_RARE_ALL_CAPS_'
        elif len(re.findall(r'[A-Z]', token[0])) == 1:
            return '_RARE_FIRST_CAP_'
        elif (token.endswith('s')):
            return  '_RARE_END_S_'
        elif len(re.findall(r'\-', token)) > 0:
            return '_RARE_HAS_DASH_'
        else:
            return  '_RARE_'


def backOffTransition(s3, s1, s2):
    if trigramTransitions.has_key((s3, s1 , s2)):
        return trigramTransitions[(s3, s1 , s2)]     
    elif bigramTransitions.has_key((s3, s2)):
        return bigramTransitions[(s3, s2)]
    elif unigramTransitions.has_key(s3):
        return unigramTransitions[s3]
    
def interpolationTransition(s3, s1, s2):
    tri = 0.0
    bi = 0.0
    uni = 0.0
    
    l1 = 0.6
    l2 = 0.3
    l3 = 0.1
    
    if trigramTransitions.has_key((s3, s1 , s2)):
        tri = math.exp(trigramTransitions[(s3, s1 , s2)]) 
    elif bigramTransitions.has_key((s3, s2)):
         bi = math.exp(bigramTransitions[(s3, s2)])
    elif unigramTransitions.has_key(s3):
        uni = math.exp(unigramTransitions[s3]);

    interpolated_prob = (l1 * tri) + (l2 * bi) + (l3 * uni)

    return math.log(interpolated_prob)


def simpleTransition(s3, s1, s2):
    # trying to compute s3|s1,s
    if trigramTransitions.has_key((s3, s1 , s2)):
        return trigramTransitions[(s3, s1 , s2)]     
    else:
        return float("-inf")  
    

def getViterbiProbabilityBigram(possible_states_per_word, words):
    pi = {}
    arg_pi = {}
    pi[(0, s2i['*'])] = 0  # 0,*
    arg_pi[(0, s2i['*'])] = []
    for k in range(1, max(possible_states_per_word.keys()) + 1):  # the words are numbered from 1 to n, 0 is special start character
        for v in possible_states_per_word[k]:
            prob2bt = {}
            for u in possible_states_per_word[k - 1]:
                q = interpolationTransition(v, None, u) #    q(v|u)
                if emissions.has_key((words[k], v)):
                    e = emissions[(words[k] , v)]   # e(words[k] | v)
                    # print words[k], '|', v , '=' , str(e)
                else:
                    e = float("-inf")
                pi_key = str(k - 1) + ',' + str(u)
                p = pi[(k - 1 , u)] + q + e
                bt = list(arg_pi[(k - 1 , u)]) 
                bt.append(u)
                prob2bt[p] = bt
            max_p = max(prob2bt.keys())
            max_bt = prob2bt[max_p]
            new_pi_key = (k, v)
            pi[new_pi_key] = max_p
            arg_pi[new_pi_key] = max_bt
    
    k = max(possible_states_per_word.keys())
    prob2bt = {}
    for v in possible_states_per_word[k]: 
        if bigramTransitions.has_key((s2i['STOP'] , v)):
            q = bigramTransitions[(s2i['STOP'] , v)]
        else:
            q = float("-inf")
            
        p = pi[(k   , v)] + q
        bt = list(arg_pi[(k, v)])             
        bt.append(v)
        prob2bt[p] = bt;  
    max_bt = prob2bt[max(prob2bt.keys())]
    max_p = max(prob2bt.keys())
    max_bt.pop(0)
    max_bt = map(lambda bt: i2s[bt], max_bt) 
    return max_bt

def getViterbiProbability(possible_states_per_word, words):
    pi = {}
    arg_pi = {}
    pi[(0, s2i['*'], s2i['*'])] = 0  # 0,*,*
    arg_pi[(0, s2i['*'], s2i['*'])] = []
    
    for k in range(1, max(possible_states_per_word.keys()) + 1):
        for v in possible_states_per_word[k]:
            for u in possible_states_per_word[k - 1]:
                prob2bt = {}
                for w in possible_states_per_word[k - 2]:
                    
                    q = interpolationTransition(v, w, u)                    
                    # print str(v) , '|' , str(w) , ',' + str(v) , '=' , str(q)
                    
                    if emissions.has_key((words[k], v)):
                        e = emissions[(words[k] , v)]
                        # print words[k], '|', v , '=' , str(e)
                    else:
                        e = float("-inf")
                   
                    pi_key = str(k - 1) + ',' + str(w) + ',' + str(u)
                    # #print 'searching pi_key', pi_key
                    p = pi[(k - 1 , w , u)] + q + e
                    bt = list(arg_pi[(k - 1, w , u)])              
                    bt.append(w)
                    prob2bt[p] = bt
                max_p = max(prob2bt.keys())
                # print 'max_p =', max_p, 'when k,v,w,u' , k, v, w, u
                max_bt = prob2bt[max_p]
                new_pi_key = (k, u , v)
                pi[new_pi_key] = max_p
                arg_pi[new_pi_key] = max_bt

    # finding max u,v
    k = max(possible_states_per_word.keys())
    prob2bt = {}
    for v in possible_states_per_word[k]:
        for u in possible_states_per_word[k - 1]:
                if trigramTransitions.has_key((s2i['STOP'], u , v)):
                    q = trigramTransitions[(s2i['STOP'], u , v)]
                else:
                    q = float("-inf") 
                p = pi[(k  , u , v)] + q
                bt = list(arg_pi[(k, u , v)])             
                bt.append(u)
                bt.append(v)
                prob2bt[p] = bt;

    max_bt = prob2bt[max(prob2bt.keys())]
    max_p = max(prob2bt.keys())
    max_bt.pop(0)
    max_bt.pop(0)
    max_bt = map(lambda bt: i2s[bt], max_bt)
    # print "sent:", words
    # print "back trace:", max_bt
    # print "probability", max_p
    # print '\n'
        
    return max_bt



def getPossibleSates(sentence):
    words = {}
    possible_states_per_word = {}
    possible_states_per_word[-1] = [s2i['*']]
    possible_states_per_word[0] = [s2i['*']]
    i = 1
    for word in sentence.split("\n"):
        word = word.strip()
        if (word == ''):
            print 'skipping empty word'
        else:
            if (seenObservations.has_key(word)):
                words[i] = word
                possible_states_per_word[i] = getWordPossibleTags(word)
            else:
                rare_word = replaceRareWord(word)
                words[i] = rare_word
                possible_states_per_word[i] = getWordPossibleTags(rare_word)
            # #print word, possible_states_per_word[i]
            i += 1
    return (possible_states_per_word, words)


    
if __name__ == "__main__":
    locale.setlocale(locale.LC_ALL, 'en_US.UTF-8')

    
    countFile = '../conll/pos-rare.count'
    testFile = '../conll/pos.test'
    outputFile = '../conll/pos.out'
    writer = open(outputFile, 'w')
    MLEstimates(countFile)
    i2s = {v:k for k, v in s2i.items()}  # create a number to state mapping for displaying the state sequence in the end
    # print s2i
    # print i2s
    testSentences = open(testFile, 'r').read().split("\n\n")
    for i, a_sentence in enumerate(testSentences):
        (possible_states_per_word, words) = getPossibleSates(a_sentence.rstrip())
        # #print a_sentence.strip().split("\n")
        if(len(words) > 0):
            # print 'line:', i
            tags = getViterbiProbability(possible_states_per_word, words)
            # tags = getViterbiProbabilityBigram(possible_states_per_word, words)
        else:
            tags = []
        original_words = a_sentence.strip().split('\n')
        for word, tag in zip(original_words, tags):
            # #print word, tag
            writer.write(word + ' ' + tag + '\n')
        if(i < len(testSentences) - 1):
            writer.write('\n')    
    writer.flush()
    writer.close()
    os.system("python eval-pos.py")
